package com.intmall.flink.operator;

import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

public class TransformPartitionTest {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        DataStreamSource<Event> stream = env.fromElements(
                new Event("Mary", "./home", 1000L),
                new Event("Bob", "./cart", 2000L),
                new Event("Alice", "./cart", 2100L),
                new Event("Bob", "./product?id=1", 3000L),
                new Event("Bob", "./product?id=2", 3500L),
                new Event("Bob", "./product?id=3", 4000L),
                new Event("Bob", "./product?id=3", 3600L),
                new Event("Alice", "./product?id=1", 5100L),
                new Event("Mary", "./center", 6000L)
        );

        // 1.随机
//        stream.shuffle().print().setParallelism(4);

        // 2.轮循
//        stream.rebalance().print().setParallelism(4);

        // 3.重缩放分区rescale
        env.addSource(new RichParallelSourceFunction<Integer>() {
            @Override
            public void run(SourceContext<Integer> sourceContext) throws Exception {
                for (int i = 1; i <= 8; i++) {
                    // 将奇数发送到索引为1的并行子任务
                    // 将偶数发送到索引为0的并行子任务
                    if (i % 2 == getRuntimeContext().getIndexOfThisSubtask()) {
                        sourceContext.collect(i);
                    }
                }
            }

            @Override
            public void cancel() {

            }
        }).setParallelism(2)
//                .rescale().print()
                .setParallelism(4);

        // 4. 广播
//        stream.broadcast().print().setParallelism(4);

        // 5. 全局
//        stream.global().print();

        // 6. 自定义分区
        env.fromElements(1, 2, 3, 4, 5, 6, 7, 8)
                .partitionCustom(new Partitioner<Integer>() {
                    @Override
                    public int partition(Integer key, int numPartitions) {
                        return key % 2;
                    }
                }, new KeySelector<Integer, Integer>() {
                    @Override
                    public Integer getKey(Integer value) throws Exception {
                        return value;
                    }
                }).print().setParallelism(4);

        env.execute();

    }
}
